# @Author     ：SEED
# @DateTime   ：2019-09-25  08:40 
# @FileName   ：fabric.py

from fabric2 import Connection
from utils.websocket_tail import Tailf
from invoke import Responder, Result


def say_yes():
    return Responder(
        pattern=r'yes/no',
        response='yes\n',
    )


class Shell(Connection):
    run_mode_remote = 'remote'
    run_mode_local = 'local'
    custom_global_env = {}

    def init_env(self, **kwargs):
        self.custom_global_env = kwargs['env']

    def run(self, command, run_mode=run_mode_remote, write=None, pty=False, webuser=None, **kwargs):
        try:
            stream_out = None
            stream_err = None
            if run_mode == self.run_mode_local:
                result = super(Shell, self).local(command, pty=pty, warn=True, out_stream=stream_out,
                                                  err_stream=stream_err,
                                                  watchers=[say_yes()],
                                                  env=self.custom_global_env, **kwargs)
            else:
                result = super(Shell, self).run(command, pty=pty, warn=True, out_stream=stream_out,
                                                err_stream=stream_err,
                                                watchers=[say_yes()],
                                                env=self.custom_global_env, **kwargs)
            exited, stdout, stderr = result.exited, result.stdout, result.stderr
            if result.failed:
                message = '[%s@%s]# %s\n[ERROR] %s' % (self.user, self.host, command, stdout + stderr)
                print(message)
            if webuser:
                if result.failed:
                    message_in = '<span style="color: #F5222D;">[%s@%s]# %s' % (self.user, self.host, command)
                    message_out = '[ERROR]' + stdout + stderr
                else:
                    message_in = '<span style="color: #52C41A;">[%s@%s]# %s</span>' % (self.user, self.host, command)
                    message_out = stdout + stderr
                # message_out = 'ERROR' + stdout + stderr
                websocket = Tailf()
                websocket.send_message(webuser, message_in)
                for m in message_out.split('\n'):
                    websocket.send_message(webuser, m)
            return result
        except Exception as e:
            message = '[%s@%s]%s' % (self.user, self.host, e)
            # print(message)
            message = '[%s@%s]# %s\n[ERROR] %s' % (self.user, self.host, command, str(e))
            if write:
                with open(write, 'a') as f:
                    f.write(message)
            elif webuser:
                message_in = '[%s@%s]# %s' % (self.user, self.host, command)
                message_out = '[ERROR] %s' % (str(e))
                Tailf.send_message(webuser, message_in)
                for m in message_out.split('\n'):
                    Tailf.send_message(webuser, m)
            result = Result(exited=-1, stderr=message, stdout=message)
            return result

    def local(self, command, write=None, webuser=None, **kwargs):
        return self.run(command, run_mode=self.run_mode_local, write=write, webuser=webuser, **kwargs)

    def get(self, remote=None, local=None, write=None, webuser=None):
        return self.sync(wtype='get', remote=remote, local=local, write=write, webuser=webuser)

    def put(self, local, remote=None, write=None, webuser=None):
        return self.sync(wtype='put', local=local, remote=remote, write=write, webuser=webuser)

    def sync(self, wtype, local=None, remote=None, write=None, webuser=None):
        try:
            if wtype == 'put':
                result = super(Shell, self).put(local, remote=remote)
                message = '[%s@%s]# [上传文件]\n[INFO] 本地:%s 上传到 [%s]:%s\n' % (
                    self.user, self.host, local, result.connection.host, remote)
                print(message)
                if write:
                    with open(write, 'a') as f:
                        f.write(message)
                elif webuser:
                    Tailf.send_message(webuser, message)
            else:
                result = super(Shell, self).get(remote, local=local)
                message = '[%s@%s]# [下载文件]\n[INFO] [%s]:%s 下载到 本地:%s\n' % (
                    self.user, self.host, result.connection.host, remote, local)
                print(message)
                if write:
                    with open(write, 'a') as f:
                        f.write(message)
                elif webuser:
                    Tailf.send_message(webuser, message)
            return result
        except Exception as e:
            if wtype == 'put':
                message = '[%s@%s]# [上传文件]\n[ERROR] [目标目录:%s][%s]\n' % (self.user, self.host, remote, str(e))
                print(message)
                if write:
                    with open(write, 'a') as f:
                        f.write(message)
                elif webuser:
                    Tailf.send_message(webuser, message)
            else:
                message = '[%s@%s]# [下载文件]\n[ERROR] [目标目录:%s][%s]\n' % (self.user, self.host, remote, str(e))
                print(message)
                if write:
                    with open(write, 'a') as f:
                        f.write(message)
                elif webuser:
                    Tailf.send_message(webuser, message)
